import math

import paddle.fluid as fluid
from paddle.fluid.initializer import MSRA
from paddle.fluid.param_attr import ParamAttr

__all__ = [
    'ShuffleNetV2'
]

from ppcls.modeling.architectures.cifar.layers import conv2d, global_avg_pool, fc


def se(x,
       reduction=4,
       name=None):
    in_channels = x.shape[1]
    channels = in_channels // reduction
    s = global_avg_pool(x)
    s = fc(s, channels, act='relu', name=name + '.squeeze')
    s = fc(s, in_channels, act='sigmoid', name=name + '.excitation')
    x = fluid.layers.elementwise_mul(x=x, y=s, axis=0)
    return x


def channel_shuffle(x, groups):
    b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
    channels_per_group = c // groups

    # reshape
    x = fluid.layers.reshape(
        x=x, shape=[b, groups, channels_per_group, h, w])

    x = fluid.layers.transpose(x=x, perm=[0, 2, 1, 3, 4])

    # flatten
    x = fluid.layers.reshape(x=x, shape=[b, c, h, w])

    return x


def basic_unit(x, use_se, name):
    c = x.shape[1] // 2
    x1, x2 = fluid.layers.split(x, num_or_sections=[c, c], dim=1)
    x2 = conv2d(x2, c, 1, bn=True, act='swish', name=name + '.conv1')
    x2 = conv2d(x2, c, 3, groups=c,
                bn=True, name=name + '.conv2')
    x2 = conv2d(x2, c, 1, bn=True, act='swish', name=name + '.conv3')
    if use_se:
        x2 = se(x2, reduction=4, name=name + ".se")
    x = fluid.layers.concat([x1, x2], axis=1)
    return channel_shuffle(x, 2)


def reduce_unit(x, out_channels, use_se, name):
    c = out_channels // 2
    # branch1
    x1 = conv2d(x, x.shape[1], 3, stride=2, groups=c,
                bn=True, name=name + ".conv4")
    x1 = conv2d(x1, c, 1, bn=True, act='swish', name=name + ".conv5")

    # branch2
    x2 = conv2d(x, c, 1, bn=True, act='swish', name=name + ".conv1")
    x2 = conv2d(x2, c, 3, stride=2, groups=c,
                bn=True, name=name + ".conv2")
    x2 = conv2d(x2, c, 1, bn=True, act='swish', name=name + ".conv3")
    if use_se:
        x2 = se(x2, reduction=4, name=name + ".se")
    x = fluid.layers.concat([x1, x2], axis=1)
    return channel_shuffle(x, 2)


def _make_layer(x, num_units, out_channels, stride, use_se, name):
    if stride == 2:
        x = reduce_unit(x, out_channels, use_se, name + ".unit1")
    else:
        x = conv2d(x, out_channels, 3, bn=True, act='swish', name=name + ".unit1")
    for i in range(1, num_units):
        x = basic_unit(x, use_se, name + f".unit{i + 1}")
    return x


class ShuffleNetV2:

    def __init__(self, stem_channels, channels_per_stage, units_per_stage, final_channels, use_se):
        self.stem_channels = stem_channels
        self.channels_per_stage = channels_per_stage
        self.units_per_stage = units_per_stage
        self.final_channels = final_channels
        self.use_se = use_se

    def net(self, input, class_dim=10):
        x = conv2d(input, self.stem_channels, 3,
                   bn=True, act='swish', name="stem")

        strides = [1, 2, 2]
        for i, c, r, s in zip(range(3), self.channels_per_stage, self.units_per_stage, strides):
            x = _make_layer(x, r, c, s, self.use_se, f"stage{i + 1}")

        x = conv2d(x, self.final_channels, 1, bn=True, act='swish', name="final_conv")
        x = global_avg_pool(x)

        x = fc(x, class_dim, name="fc")
        return x
